package com.ruoyi.video.server.utils;

import com.ruoyi.video.server.ws.RtspWebSocket;
import lombok.extern.slf4j.Slf4j;
import org.apache.commons.io.IOUtils;
import org.bytedeco.ffmpeg.avcodec.AVPacket;
import org.bytedeco.ffmpeg.global.avcodec;
import org.bytedeco.javacv.*;

import java.io.ByteArrayOutputStream;

@Slf4j
public class RtspThread extends Thread{
    private String rtsp;
    private Boolean runing;
    private Boolean isH264 = false;

    /**
     * 读流器
     */
    private FFmpegFrameGrabber grabber;
    /**
     * 转码器
     */
    private FFmpegFrameRecorder recorder;
    /**
     * 转FLV格式的头信息<br/>
     * 如果有第二个客户端播放首先要返回头信息
     */
    private byte[] headers;
    /**
     * 保存转换好的流
     */
    private ByteArrayOutputStream stream;

    public RtspThread(String rtsp){
        this.rtsp = rtsp;
        this.runing = true;
    }

    public String getRtsp() {
        return rtsp;
    }

    public void setRtsp(String rtsp) {
        this.rtsp = rtsp;
    }

    public Boolean getRuning() {
        return runing;
    }

    public void setRuning(Boolean runing) {
        this.runing = runing;
    }

    public byte[] getHeaders() {
        return headers;
    }

    public void setHeaders(byte[] headers) {
        this.headers = headers;
    }

    @Override
    public void run() {
        try {
            grabber = new FFmpegFrameGrabber(rtsp);
            if ("rtsp".equals(rtsp.substring(0, 4))) {
                grabber.setOption("rtsp_transport", "tcp");
                grabber.setOption("stimeout", "5000000");
            }
            grabber.start();
            if (avcodec.AV_CODEC_ID_H264 == grabber.getVideoCodec()
                    && (grabber.getAudioChannels() == 0 || avcodec.AV_CODEC_ID_AAC == grabber.getAudioCodec())) {
                log.info("摄像机:{}是H264编码", rtsp);
                isH264 = true;
                doH264Stream();
            } else {
                log.info("摄像机:{}不是H264，需要转码", rtsp);
                doOtherStream();
            }
        } catch (Exception e) {
            log.error(e.getMessage(), e);
        } finally {
            try {
                IOUtils.close(grabber);
                IOUtils.close(recorder);
                IOUtils.close(stream);
            }catch (Exception e){
                e.printStackTrace();
            }
            log.info("摄像机:{}的当前转流已经结束", rtsp);
        }
    }

    public void doH264Stream() throws FrameRecorder.Exception, FrameGrabber.Exception, InterruptedException {
        // 来源视频H264格式,音频AAC格式
        // 无须转码，更低的资源消耗，更低的延迟
        stream = new ByteArrayOutputStream();
        recorder = new FFmpegFrameRecorder(stream, grabber.getImageWidth(), grabber.getImageHeight(),
                grabber.getAudioChannels());
        recorder.setInterleaved(true);
        recorder.setVideoOption("preset", "ultrafast");
        recorder.setVideoOption("tune", "zerolatency");
        recorder.setVideoOption("crf", "25");
        recorder.setFrameRate(grabber.getFrameRate());
        recorder.setSampleRate(grabber.getSampleRate());
        if (grabber.getAudioChannels() > 0) {
            recorder.setAudioChannels(grabber.getAudioChannels());
            recorder.setAudioBitrate(grabber.getAudioBitrate());
            recorder.setAudioCodec(grabber.getAudioCodec());
        }
        recorder.setFormat("flv");
        recorder.setVideoBitrate(grabber.getVideoBitrate());
        recorder.setVideoCodec(grabber.getVideoCodec());
        recorder.start(grabber.getFormatContext());
        if (headers == null) {
            headers = stream.toByteArray();
            stream.reset();
            writeResponse(headers);
        }
        int nullNumber = 0;
        while (runing) {
            AVPacket k = grabber.grabPacket();
            if (k != null) {
                try {
                    recorder.recordPacket(k);
                } catch (Exception e) {
                }
                if (stream.size() > 0) {
                    byte[] b = stream.toByteArray();
                    stream.reset();
                    writeResponse(b);
                }
                avcodec.av_packet_unref(k);
            } else {
                nullNumber++;
                if (nullNumber > 200) {
                    break;
                }
            }
            sleep(5);
        }
    }

    public void doOtherStream() throws FrameRecorder.Exception, FrameGrabber.Exception, InterruptedException {
        log.info("this url:{} converterTranFactories start", "dd");
        grabber.setFrameRate(25);
        if (grabber.getImageWidth() > 1920) {
            grabber.setImageWidth(1920);
        }
        if (grabber.getImageHeight() > 1080) {
            grabber.setImageHeight(1080);
        }
        stream = new ByteArrayOutputStream();
        recorder = new FFmpegFrameRecorder(stream, grabber.getImageWidth(), grabber.getImageHeight(),
                grabber.getAudioChannels());
        recorder.setInterleaved(true);
        recorder.setVideoOption("preset", "ultrafast");
        recorder.setVideoOption("tune", "zerolatency");
        recorder.setVideoOption("crf", "25");
        recorder.setGopSize(50);
        recorder.setFrameRate(25);
        recorder.setSampleRate(grabber.getSampleRate());
        if (grabber.getAudioChannels() > 0) {
            recorder.setAudioChannels(grabber.getAudioChannels());
            recorder.setAudioBitrate(grabber.getAudioBitrate());
            recorder.setAudioCodec(avcodec.AV_CODEC_ID_AAC);
        }
        recorder.setFormat("flv");
        recorder.setVideoBitrate(grabber.getVideoBitrate());
        recorder.setVideoCodec(avcodec.AV_CODEC_ID_H264);
        recorder.start();
        if (headers == null) {
            headers = stream.toByteArray();
            stream.reset();
            writeResponse(headers);
        }
        int nullNumber = 0;
        while (runing) {
            // 抓取一帧
            Frame f = grabber.grab();
            if (f != null) {
                try {
                    // 转码
                    recorder.record(f);
                } catch (Exception e) {
                }
                if (stream.size() > 0) {
                    byte[] b = stream.toByteArray();
                    stream.reset();
                    writeResponse(b);
                }
            } else {
                nullNumber++;
                if (nullNumber > 200) {
                    break;
                }
            }
            sleep(5);
        }
    }
    /**
     * 输出FLV视频流
     *
     * @param b
     */
    public void writeResponse(byte[] b) {
        RtspWebSocket.sendData(rtsp, b);
    }
}
